from asyncio import as_completed
import os
import sys
import hashlib
from tqdm import tqdm
from functools import partial
from concurrent.futures import ThreadPoolExecutor, wait, as_completed

try:
    from .cvio import cvio
except Exception as e1:
    try:
        from cvtools import cvio
    except Exception as e2:
        print(e1, e2)
        from cvio import cvio


def compute_mdfile(file):
    with open(file, 'rb') as fp:
        data = fp.read()
    return hashlib.md5(data).hexdigest()


def compute_mdfile_future(file, results=[]):
    md5 = compute_mdfile(file)
    results.append((file, md5))


def compute_mdfiles(src, ext_type=None, recursive=False, num_workers=4, silent=False):
    pool = ThreadPoolExecutor(num_workers)

    if ext_type == None:
        load = partial(cvio.load_image_list,
                       recursive=recursive, silent=silent)
    elif isinstance(ext_type, str):
        load = partial(cvio.load_ext_list, ext_type=ext_type,
                       recursive=recursive, silent=silent)
    else:
        raise TypeError

    files = load(src)
    results = []
    n = len(files)
    tasks = []
    for i, file in enumerate(files, 1):
        desc = '[%d/%d][%.2f%%] %s\r' % (i, n, i / n * 100, os.path.basename(file))

        def callback(*args, desc=desc):
            sys.stdout.write(desc)
            sys.stdout.flush()
        t = pool.submit(compute_mdfile_future, file, results)
        t.add_done_callback(callback)
        tasks.append(t)
    for p in as_completed(tasks):
        try:
            p.result()
        except Exception as e:
            print(e)
            pool.shutdown(wait=True)
    sys.stdout.write('finished\n')
    return results

def save_mdfiles(results, dst):
    with open(dst, 'w') as fp:
        for file, md5 in results:
            fp.write('%s %s\n' % (file, md5))

if __name__ == '__main__':
    # src = r'G:\data\datasets\drink\daiding_101\proj\train'
    # src = 'coco(5)'
    src = r'G:\data\datasets\drink\daiding_101\proj\train\repeats'
    results = compute_mdfiles(src, num_workers=32)
    save_mdfiles(results, 'repeats/mdfiles.txt')
